#include "KdTree.h"

#include <algorithm>

#include "Agent.h"
#include "Definitions.h"
#include "RVOSimulator.h"

namespace RVO
{
	const size_t RVO_MAX_LEAF_SIZE = 10; // 定义一个常量，表示KD树叶节点的最大大小

	KdTree::KdTree(RVOSimulator *sim) : sim_(sim) {} // 定义KD树的构造函数，接受一个指向RVOSimulator对象的指针，并将其初始化为类成员sim_

	// 用于构建智能体的KD树。KD树是一个用于高效查询和邻居搜索的数据结构，常用于多智能体系统中的避碰计算。
	void KdTree::buildAgentTree()
	{
		agents_ = sim_->agents_; // 将仿真器中的智能体列表复制到KD树的智能体列表中

		if (!agents_.empty())
		{
			agentTree_.resize(2 * agents_.size() - 1);     // 调整KD树节点列表的大小，确保足够存储所有节点
			buildAgentTreeRecursive(0, agents_.size(), 0); // 递归构建KD树，从根节点开始
		}
	}

	// 用于递归地构建智能体的KD树。
	void KdTree::buildAgentTreeRecursive(size_t begin, size_t end, size_t node)
	{
		agentTree_[node].begin = begin;			       // 设置当前节点处理的智能体范围的起始索引
		agentTree_[node].end = end;			       // 设置当前节点处理的智能体范围的结束索引
		agentTree_[node].minCoord = agents_[begin]->position_; // 初始化当前节点的最小坐标
		agentTree_[node].maxCoord = agents_[begin]->position_; // 初始化当前节点的最大坐标

		for (size_t i = begin + 1; i < end; ++i)
		{
			agentTree_[node].maxCoord[0] = std::max(agentTree_[node].maxCoord[0], agents_[i]->position_.x()); // 更新最大x坐标
			agentTree_[node].minCoord[0] = std::min(agentTree_[node].minCoord[0], agents_[i]->position_.x()); // 更新最小x坐标
			agentTree_[node].maxCoord[1] = std::max(agentTree_[node].maxCoord[1], agents_[i]->position_.y()); // 更新最大y坐标
			agentTree_[node].minCoord[1] = std::min(agentTree_[node].minCoord[1], agents_[i]->position_.y()); // 更新最小y坐标
			agentTree_[node].maxCoord[2] = std::max(agentTree_[node].maxCoord[2], agents_[i]->position_.z()); // 更新最大z坐标
			agentTree_[node].minCoord[2] = std::min(agentTree_[node].minCoord[2], agents_[i]->position_.z()); // 更新最小z坐标
		}

		if (end - begin > RVO_MAX_LEAF_SIZE)
		{
			/* 当前节点不是叶节点。 */
			size_t coord;

			// 选择划分维度
			if (agentTree_[node].maxCoord[0] - agentTree_[node].minCoord[0] > agentTree_[node].maxCoord[1] - agentTree_[node].minCoord[1] && agentTree_[node].maxCoord[0] - agentTree_[node].minCoord[0] > agentTree_[node].maxCoord[2] - agentTree_[node].minCoord[2])
			{
				coord = 0; // 按x轴划分
			}
			else if (agentTree_[node].maxCoord[1] - agentTree_[node].minCoord[1] > agentTree_[node].maxCoord[2] - agentTree_[node].minCoord[2])
			{
				coord = 1; // 按y轴划分
			}
			else
			{
				coord = 2; // 按z轴划分
			}

			const float splitValue = 0.5f * (agentTree_[node].maxCoord[coord] + agentTree_[node].minCoord[coord]); // 计算划分值

			size_t left = begin;
			size_t right = end;

			while (left < right)
			{
				while (left < right && agents_[left]->position_[coord] < splitValue)
				{
					++left; // 移动左指针直到找到需要交换的元素
				}

				while (right > left && agents_[right - 1]->position_[coord] >= splitValue)
				{
					--right; // 移动右指针直到找到需要交换的元素
				}

				if (left < right)
				{
					std::swap(agents_[left], agents_[right - 1]); // 交换两个元素
					++left;					      // 左指针向右移动
					--right;				      // 右指针向左移动
				}
			}

			size_t leftSize = left - begin;

			if (leftSize == 0)
			{
				++leftSize;
				++left;
				++right;
			}

			agentTree_[node].left = node + 1;	      // 设置左子节点索引
			agentTree_[node].right = node + 2 * leftSize; // 设置右子节点索引

			buildAgentTreeRecursive(begin, left, agentTree_[node].left); // 递归构建左子树
			buildAgentTreeRecursive(left, end, agentTree_[node].right);  // 递归构建右子树
		}
	}

	// 用于计算给定智能体的邻居智能体。通过调用递归查询函数，遍历KD树以查找在给定范围内的邻居。
	void KdTree::computeAgentNeighbors(Agent *agent, float rangeSq) const
	{
		queryAgentTreeRecursive(agent, rangeSq, 0); // 调用递归查询函数，从根节点开始查找邻居
	}

	// 用于递归地查询KD树，以查找在给定范围内的邻居智能体。
	void KdTree::queryAgentTreeRecursive(Agent *agent, float &rangeSq, size_t node) const
	{
		if (agentTree_[node].end - agentTree_[node].begin <= RVO_MAX_LEAF_SIZE)
		{
			// 当前节点是叶节点，遍历该节点中的所有智能体，检查是否在范围内
			for (size_t i = agentTree_[node].begin; i < agentTree_[node].end; ++i)
			{
				agent->insertAgentNeighbor(agents_[i], rangeSq);
			}
		}
		else
		{
			// 当前节点是内部节点，计算目标智能体到左右子节点的距离平方
			const float distSqLeft =
			    sqr(std::max(0.0f, agentTree_[agentTree_[node].left].minCoord[0] - agent->position_.x())) +
			    sqr(std::max(0.0f, agent->position_.x() - agentTree_[agentTree_[node].left].maxCoord[0])) +
			    sqr(std::max(0.0f, agentTree_[agentTree_[node].left].minCoord[1] - agent->position_.y())) +
			    sqr(std::max(0.0f, agent->position_.y() - agentTree_[agentTree_[node].left].maxCoord[1])) +
			    sqr(std::max(0.0f, agentTree_[agentTree_[node].left].minCoord[2] - agent->position_.z())) +
			    sqr(std::max(0.0f, agent->position_.z() - agentTree_[agentTree_[node].left].maxCoord[2]));

			const float distSqRight =
			    sqr(std::max(0.0f, agentTree_[agentTree_[node].right].minCoord[0] - agent->position_.x())) +
			    sqr(std::max(0.0f, agent->position_.x() - agentTree_[agentTree_[node].right].maxCoord[0])) +
			    sqr(std::max(0.0f, agentTree_[agentTree_[node].right].minCoord[1] - agent->position_.y())) +
			    sqr(std::max(0.0f, agent->position_.y() - agentTree_[agentTree_[node].right].maxCoord[1])) +
			    sqr(std::max(0.0f, agentTree_[agentTree_[node].right].minCoord[2] - agent->position_.z())) +
			    sqr(std::max(0.0f, agent->position_.z() - agentTree_[agentTree_[node].right].maxCoord[2]));

			// 优先查询距离较小的子节点
			if (distSqLeft < distSqRight)
			{
				if (distSqLeft < rangeSq)
				{
					queryAgentTreeRecursive(agent, rangeSq, agentTree_[node].left);

					if (distSqRight < rangeSq)
					{
						queryAgentTreeRecursive(agent, rangeSq, agentTree_[node].right);
					}
				}
			}
			else
			{
				if (distSqRight < rangeSq)
				{
					queryAgentTreeRecursive(agent, rangeSq, agentTree_[node].right);

					if (distSqLeft < rangeSq)
					{
						queryAgentTreeRecursive(agent, rangeSq, agentTree_[node].left);
					}
				}
			}
		}
	}
}
